import argparse
from typing import Dict, Any

import ray
import yaml
import json
from modules import MFCCPredictor
# from storage import Boto3Storage
from utils.tools import notify, parse_s3url
from utils.logger import get_logger
from ray.data import DataContext

logger = get_logger()

DataContext.get_current().enable_progress_bars = False

parser = argparse.ArgumentParser(
    prog='embedding',
    description='audio embedding')

parser.add_argument("--inputs",
                    "-i",
                    type=str,
                    help="json argument")
parser.add_argument("--config_path",
                    "-c",
                    type=str,
                    help="json argument")
args = parser.parse_args()

with open(args.config_path) as fp:
    config = yaml.safe_load(fp)


request_data = json.loads(args.inputs)
callback_url = request_data.get("callback_url")
task_id = request_data.get("task_id")
tasks = request_data.get("tasks")
# storage_cfg = config["storage"]
# bucket_name = storage_cfg.get("bucket_name")
# storage = Boto3Storage(storage_cfg['addr'],
#                        storage_cfg['accessKey'],
#                        storage_cfg['secretKey'])

def set_progress_status(row: Dict[str, Any]) -> Dict[str, Any]:
    row["task_id"] = task_id
    row["callback_url"] = callback_url
    row["embedding_progress"] = "pending"
    return row

# notify(tasks, callback_url, task_id, "pending")

ray_cfg = config['model']['ray_options']
results = []
for row in (
        ray.data.from_items(tasks)
                .map(set_progress_status)
                .map(MFCCPredictor,
                     concurrency=ray_cfg['concurrency'],
                     num_cpus=ray_cfg['num_cpus'],
                     fn_constructor_kwargs={"cfg": config},
                     # runtime_env={"env_vars": {
                     #     "LD_LIBRARY_PATH": "/root/miniconda3/envs/auto-match/lib:$LD_LIBRARY_PATH"}}
                     )
                .iter_rows()
):
    # bucket_name, obj_name = parse_s3url(row['meta_path'])
    # try:
    #     storage.put_json(row, row['meta_path'])
    #     logger.info(f"upload meta, meta path is {row['meta_path']}")
    # except:
    #     row["embedding_progress"] = "failed"

    # del row["silent_points"]
    logger.info(row)
    # notify([row], callback_url, task_id, "part_finished")
    results.append(row)


# notify(results, callback_url, task_id, "all_finished")
